﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace DecisionTree
{
    class DecisionNode
    {
        private int[] indexes;

        public int IndexCount { get; private set; }
        public DecisionTree Owner { get; private set; }
        public DecisionNode Parent { get; private set; }
        public int Level { get; private set; }

        public DecisionNode[] Branches { get; set; }
        
        public int? Output;
        public double[] OutputHistogram;
        public int SplitFeature;

        public bool IsRoot
        {
            get { return Parent == null; }
        }

        public bool IsLeaf
        {
            get { return Branches == null || Branches.Length == 0; }
        }

        public DecisionNode(DecisionTree owner, DecisionNode parent, int[] indexes)
        {
            this.indexes = indexes;
            this.Parent = parent;
            this.Owner = owner;
            this.IndexCount = indexes.Length;

            this.Level = IsRoot ? 0 : parent.Level + 1;
        }

        private int[] GetAttributeCandidates()
        {
            int[] rVal = new int[Parameters.UsedFeatureCountAtEveryNode];

            for (int i = 0; i < rVal.Length; i++)
                rVal[i] = Owner.Owner.rnd.Next(28 * 28);

            return rVal;
        }

        public void Split()
        {
            int[] attributeCandidates;
            int[][][] partitions;
            int bestAttributeIndex;
            double bestScore;

            // If all examples are for the same class, return the single-node
            // tree with the output label corresponding to this common class.
            double entropy = Tools.Entropy(this.Owner.outputs, this.indexes, Parameters.NumOutputClasses);

            if (entropy == 0)
            {
                this.Output = this.Owner.outputs[this.indexes[0]];
                this.OutputHistogram = Tools.GetHistogram(this.Owner.outputs, this.indexes, Parameters.NumOutputClasses, this.Level);
                return;
            }

            if (entropy < Parameters.EntropyTermination || this.Level > Parameters.MaxDepth || this.IndexCount < Parameters.MinNumSamplesPerLeaf)
            {
                this.Output = Tools.MostCommon(this.Owner.outputs, this.indexes, Parameters.NumOutputClasses);
                this.OutputHistogram = Tools.GetHistogram(this.Owner.outputs, this.indexes, Parameters.NumOutputClasses, this.Level);
                return;
            }          

            attributeCandidates = GetAttributeCandidates();

            double[] scores = new double[attributeCandidates.Length];
            partitions = new int[attributeCandidates.Length][][];

            for (int i = 0; i < attributeCandidates.Length; i++)
            {
                scores[i] = ComputeGainRatio(Owner.inputs, Owner.outputs, this.indexes, attributeCandidates[i], out partitions[i], entropy);
            }

            bestAttributeIndex = scores.MaxIndex();
            bestScore = scores[bestAttributeIndex];

            if (!(bestScore > 0))
            {
                this.Output = Tools.MostCommon(this.Owner.outputs, this.indexes, Parameters.NumOutputClasses);
                this.OutputHistogram = Tools.GetHistogram(this.Owner.outputs, this.indexes, Parameters.NumOutputClasses, this.Level);
                return;
            }
            
            this.SplitFeature = attributeCandidates[bestAttributeIndex];

            this.Branches = new DecisionNode[2];

            this.Branches[0] = new DecisionNode(this.Owner, this, partitions[bestAttributeIndex][0]);
            this.Branches[1] = new DecisionNode(this.Owner, this, partitions[bestAttributeIndex][1]);
        }

        public static double ComputeGainRatio(byte[][] inputs, byte[] outputs, int[] indexes, int attribute, out int[][] partitions, double currentEntropy)
        { 
            double infoGain = currentEntropy - ComputeInfo(inputs, outputs, indexes, attribute, out partitions);
            double splitInfo = Tools.SplitInformation(indexes.Length, partitions);

            double gainRatio = splitInfo == 0 ? 0 : infoGain / splitInfo;

            return gainRatio;
        }

        public static double ComputeInfo(byte[][] inputs, byte[] outputs, int[] indexes, int attribute, out int[][] partitions)
        {
            partitions = new int[2][];

            List<int> tempPartitions0 = new List<int>(indexes.Length);
            List<int> tempPartitions1 = new List<int>(indexes.Length);            
            
            for (int i = 0; i < indexes.Length; i++)
            {
                if(inputs[indexes[i]][attribute] == 0)
                    tempPartitions0.Add(indexes[i]);
                else
                    tempPartitions1.Add(indexes[i]);
            }

            partitions[0] = tempPartitions0.ToArray();
            partitions[1] = tempPartitions1.ToArray();

            // Compute the information gain obtained by using
            // this current attribute as the next decision node.
            double info = 0;

            for (int i = 0; i < partitions.Length; i++)
            {
                // Check the entropy gain originating from this partitioning
                double e = Tools.Entropy(outputs, partitions[i], Parameters.NumOutputClasses);

                info += ((double)partitions[i].Length / indexes.Length) * e;
            }

            return info;
        }

        public bool IsValid(byte[] p)
        {
            if (p[SplitFeature] == 0)
                return true;
            return false;
        }

        public LightDecisionNode ToLightDecisionNode(LightDecisionNode parent)
        {
            LightDecisionNode rVal = new LightDecisionNode(parent);

            LightDecisionNode[] branches = null;

            if ( !( this.Branches == null || this.Branches.Length == 0) )
            {
                branches = new LightDecisionNode[2];
                for (int i = 0; i < this.Branches.Length; i++)
                {
                    branches[i] = this.Branches[i].ToLightDecisionNode(rVal);
                }
            }

            rVal.SetValues(branches, this.IndexCount, this.Level, this.SplitFeature, this.Output, this.OutputHistogram);

            return rVal;
        }
    }
}
